import socket
import select
from threading import Thread

# Protocol Constants
ADDR = "239.192.152.143"
PORT = 6771
BTSEARCH  = "BT-SEARCH * HTTP/1.1\r\nHost: 239.192.152.143:6771\r\nPort: %d\r\nInfohash: %s\r\n\r\n\r\n"

class MulticastUDPSocket(socket.socket):
    """ Multicast UDP Socket """
    def __init__(self, address = ADDR, port = PORT):       
        socket.socket.__init__(self, socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        self._s_address = address
        self._s_port = port

        try:
            self.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        except:
            pass
        self.setsockopt(socket.SOL_IP, socket.IP_MULTICAST_LOOP, 1)
        self.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 255)
        self.bind(("", port))
        self.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, socket.inet_aton(address) + socket.inet_aton("0.0.0.0"))

    def close(self):
        self.setsockopt(socket.SOL_IP, socket.IP_DROP_MEMBERSHIP, socket.inet_aton(self._s_address) + socket.inet_aton('0.0.0.0'))
        socket.socket.close(self)
        
class LSD(Thread):
    """ Local Service Discovery """
    def __init__(self, callback = None):
        """ Initialize LSD """
        Thread.__init__(self)
        
        self.active = False
        self.callback = callback
        self.sock = None
                
    def run(self):
        """ Start Listening """
        self.sock = MulticastUDPSocket()
        self.active = True
        while self.active:
            rr, ww, ee = select.select([self.sock], [], [], 10)
            for r in rr:
                sender, data = r.recvfrom(256)
                self.on_announce(sender, data)
        self.sock.close()
        
    def on_announce(self, sender, data):
        """ Got an announce """
        print sender, data
        #self.callback(ip, port, ih)
            
    def announce(self, infohash, listen_port):
        """ Send an announce """
        if not self.sock or not self.active:
            return True # no retry
        
        btsearch  = BTSEARCH % (listen_port, infohash)
        try:
            sent = self.sock.sendto(btsearch, 0, (self.sock._s_address, self.sock._s_port))
        except:
            success = False
        else:
            success = sent == len(btsearch)
            
        return success
        
    def close(self):
        """ Close LSD """
        self.active = False
        
if __name__ == "__main__":
    import time
    l = LSD()
    l.start()
    print l.announce("01234567890123456789", 49152)
    time.sleep(5)
    l.close()
